
#include "DCTLee.cuh"
#include "DCTFFT.cuh"
#include "DCTTest.h"
#include "dct.h"

using namespace fabric_space;

void PrintVec(const std::vector<float> &vec, std::string fn)
{
    std::ofstream ofs(fn.c_str(), std::ios_base::trunc | std::ios_base::out);
    for (size_t i = 0; i < vec.size(); ++i)
    {
        ofs << vec[i] << std::endl;
    }
    ofs.close();
}

void DCTTest::TestCPUDCT(const std::vector<float> &data, unsigned int dim)
{
    // PrintVec(data, "data_cpu");
    auto data2 = data;
    transpose(data.data(), data2.data(), dim, dim);
    for (unsigned int r = 0; r < dim; ++r)
    {
        for (unsigned int c = 0; c < dim; ++c)
        {
            assert(data[r * dim + c] == data2[c * dim + r]);
        }
    }
    std::vector<float> cos_dct(dim), cos_idct(dim);
    precompute_dct_cos(cos_dct.data(), dim);
    precompute_idct_cos(cos_idct.data(), dim);
    // PrintVec(cos_dct, "cos_dct_cpu");
    // PrintVec(cos_idct, "cos_idct_cpu");
    std::vector<float> out1(dim * dim), out2(dim * dim), buf(dim * dim);
    dct((float *)data.data(), out1.data(), buf.data(), cos_dct.data(), dim, dim, 8);
    // PrintVec(out1, "out1_dct_cpu");
    for (unsigned i = 0; i < dim * dim; ++i)
    {
        out1[i] *= (float)(2.0 / dim);
    }
    // PrintVec(out1, "out1_dct_2_cpu");
    clock_t start_time, end_time;
    start_time = clock();
    dct2d((float *)data.data(), out1.data(), buf.data(), cos_dct.data(), cos_dct.data(), dim, dim, 8);
    // PrintVec(out1, "out1_cpu");
    end_time = clock();
    std::cout << "CPU dct2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    start_time = clock();
    idct2d(out1.data(), out2.data(), buf.data(), cos_idct.data(), cos_idct.data(), dim, dim, 8);
    // PrintVec(out2, "out2_cpu");
    end_time = clock();
    std::cout << "CPU idct2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    double sum1 = 0.0, sum2 = 0.0;
    for (unsigned int r = 0; r < dim; ++r)
    {
        for (unsigned int c = 0; c < dim; ++c)
        {
            double dr = (data[r * dim + c] - out2[r * dim + c]);
            sum1 += dr * dr;
            sum2 += data[r * dim + c] * out2[r * dim + c];
        }
    }
    std::cout << "DCT/IDCT error:" << std::sqrt(sum1 / sum2) << std::endl;
    start_time = clock();
    dst2d((float *)data.data(), out1.data(), buf.data(), cos_dct.data(), cos_dct.data(), dim, dim, 8);
    end_time = clock();
    std::cout << "CPU dst2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    start_time = clock();
    idst2d(out1.data(), out2.data(), buf.data(), cos_idct.data(), cos_idct.data(), dim, dim, 8);
    end_time = clock();
    std::cout << "CPU idst2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    sum1 = 0.0;
    sum2 = 0.0;
    for (unsigned int r = 0; r < dim; ++r)
    {
        for (unsigned int c = 0; c < dim; ++c)
        {
            double dr = (data[r * dim + c] - out2[r * dim + c]);
            sum1 += dr * dr;
            sum2 += data[r * dim + c] * out2[r * dim + c];
        }
    }
    std::cout << "DST/IDST error:" << std::sqrt(sum1 / sum2) << std::endl;
}

/**
 * Find and set CUDA device, return false if no CUDA device is found in system.
 * @author zpfeng
 */
bool FindAndSetCUDADevice(int dev_id)
{
    bool res = __cudaSafeCallCheck(cudaSetDevice(dev_id));
    int major = 0, minor = 0;
    res = res && __cudaSafeCallCheck(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev_id));
    res = res && __cudaSafeCallCheck(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev_id));
    if (res)
    {
        printf("\nGPU Device %d with compute capability %d.%d\n", dev_id, major, minor);
        return true;
    }
    else
    {
        return false;
    }
}

void DCTTest::TestGPUDCT(const std::vector<float> &data, unsigned int dim)
{
    if (!FindAndSetCUDADevice(0))
    {
        return;
    }
    // test Lee DCT
    std::vector<float> cos_dct(dim), cos_idct(dim);
    dct_op::PrecomputeDCTCos(cos_dct.data(), dim);
    dct_op::PrecomputeIDCTCos(cos_idct.data(), dim);
    std::vector<float> out1(dim * dim), out2(dim * dim), buf(dim * dim);
    float *g_data = nullptr, *g_out1 = nullptr, *g_out2 = nullptr, *g_buf = nullptr, *g_cos_dct = nullptr, *g_cos_idct = nullptr;
    AllocateCopyCUDA(&g_data, data.data(), data.size());
    AllocateCopyCUDA(&g_out1, out1.data(), out1.size());
    AllocateCopyCUDA(&g_out2, out2.data(), out2.size());
    AllocateCopyCUDA(&g_buf, buf.data(), buf.size());
    AllocateCopyCUDA(&g_cos_dct, cos_dct.data(), cos_dct.size());
    AllocateCopyCUDA(&g_cos_idct, cos_idct.data(), cos_idct.size());
    // PrintCUDAArrayToFile(g_cos_dct, cos_dct.size(), "cos_dct_gpu");
    // PrintCUDAArrayToFile(g_cos_idct, cos_idct.size(), "cos_idct_gpu");
    dct_op::MatTransposeV2(g_data, g_buf, dim, dim);
    dct_op::MatTransposeV2(g_buf, g_out1, dim, dim);
    double sum1 = 0.0, sum2 = 0.0;
    MemcpyDeviceToHost(out2.data(), g_out1, out2.size());
    for (unsigned int r = 0; r < dim; ++r)
    {
        for (unsigned int c = 0; c < dim; ++c)
        {
            double dr = (data[r * dim + c] - out2[r * dim + c]);
            sum1 += dr * dr;
            sum2 += data[r * dim + c] * out2[r * dim + c];
        }
    }
    std::cout << "transpose error: " << std::sqrt(sum1 / sum2) << std::endl;
    // PrintCUDAArrayToFile(g_data, data.size(), "data_gpu");
    dct_op::DCTKernel<<<CeilDiv(dim, (unsigned)256), min((unsigned)256, dim)>>>(g_data, g_out1, g_buf, g_cos_dct, dim, dim);
    cudaDeviceSynchronize();
    // PrintCUDAArrayToFile(g_out1, out1.size(), "out1_dct_gpu");
    dct_op::MultiplyScaleCoefKernel<<<CeilDiv(dim * dim, (unsigned)256), min((unsigned)256, dim * dim)>>>(g_out1, dim, dim);
    cudaDeviceSynchronize();
    // PrintCUDAArrayToFile(g_out1, out1.size(), "out1_dct_2_gpu");
    clock_t start_time, end_time;
    start_time = clock();
    dct_op::DCT2D(g_data, g_out1, g_buf, g_cos_dct, g_cos_dct, dim, dim);
    // PrintCUDAArrayToFile(g_out1, out1.size(), "out1_gpu");
    end_time = clock();
    std::cout << "GPU dct2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    start_time = clock();
    dct_op::IDCT2D(g_out1, g_out2, g_buf, g_cos_idct, g_cos_idct, dim, dim);
    // PrintCUDAArrayToFile(g_out2, out2.size(), "out2_gpu");
    end_time = clock();
    std::cout << "GPU idct2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    MemcpyDeviceToHost(out2.data(), g_out2, out2.size());
    for (unsigned int r = 0; r < dim; ++r)
    {
        for (unsigned int c = 0; c < dim; ++c)
        {
            double dr = (data[r * dim + c] - out2[r * dim + c]);
            sum1 += dr * dr;
            sum2 += data[r * dim + c] * out2[r * dim + c];
        }
    }
    std::cout << "DCT/IDCT error: " << std::sqrt(sum1 / sum2) << std::endl;
    start_time = clock();
    dct_op::DST2D(g_data, g_out1, g_buf, g_cos_dct, g_cos_dct, dim, dim);
    end_time = clock();
    std::cout << "GPU dst2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    start_time = clock();
    dct_op::IDST2D(g_out1, g_out2, g_buf, g_cos_idct, g_cos_idct, dim, dim);
    end_time = clock();
    std::cout << "GPU idst2d time: " << (double)(end_time - start_time) / CLOCKS_PER_SEC << std::endl;
    sum1 = 0.0;
    sum2 = 0.0;
    for (unsigned int r = 0; r < dim; ++r)
    {
        for (unsigned int c = 0; c < dim; ++c)
        {
            double dr = (data[r * dim + c] - out2[r * dim + c]);
            sum1 += dr * dr;
            sum2 += data[r * dim + c] * out2[r * dim + c];
        }
    }
    std::cout << "DST/IDST error: " << std::sqrt(sum1 / sum2) << std::endl;
    DestroyCUDA(g_data);
    DestroyCUDA(g_out1);
    DestroyCUDA(g_out2);
    DestroyCUDA(g_buf);
    DestroyCUDA(g_cos_dct);
    DestroyCUDA(g_cos_idct);
}

void DCTTest::TestGPUDCTFFT(const std::vector<float> &data, unsigned int dim)
{
    /*unsigned int mat_size = _row * _col;
    std::vector<float> test_mat(mat_size, 0.0f), test_mat2(mat_size, 0.0f);
    for (unsigned int r = 0; r < _row; ++r)
    {
        for (unsigned int c = 0; c < _col; ++c)
        {
            test_mat[r * _col + c] = ((float)std::rand() / RAND_MAX) - 0.5f;
        }
    }
    float *test_mat_gpu = nullptr, *test_mat_gpu2 = nullptr;
    AllocateCopyCUDA(&test_mat_gpu, test_mat.data(), mat_size);
    AllocateMemsetCUDA(&test_mat_gpu2, 0, mat_size);
    DCT2D(test_mat_gpu, test_mat_gpu2);
    // PrintCUDA2DArray(test_mat_gpu, _col, _row, "test_mat_gpu_before_idct2d");
    // PrintCUDA2DArray(test_mat_gpu2, _col, _row, "test_mat_gpu2_before_idct2d");
    IDCT2D(test_mat_gpu2, test_mat_gpu);
    // PrintCUDA2DArray(test_mat_gpu, _col, _row, "test_mat_gpu_after_idct2d");
    // PrintCUDA2DArray(test_mat_gpu2, _col, _row, "test_mat_gpu2_after_idct2d");
    MemcpyDeviceToHost(test_mat2.data(), test_mat_gpu, mat_size);
    float sum1 = 0.0f, sum2 = 0.0f;
    for (unsigned int r = 0; r < _row; ++r)
    {
        for (unsigned int c = 0; c < _col; ++c)
        {
            float dr = test_mat[r * _col + c] - test_mat2[r * _col + c];
            sum1 += dr * dr;
            sum2 += test_mat[r * _col + c] * test_mat2[r * _col + c];
        }
    }
    std::cout << "DCT2D/IDCT2D verify error: " << std::sqrt(sum1 / sum2) << std::endl;
    DestroyCUDA(test_mat_gpu);
    DestroyCUDA(test_mat_gpu2);*/
}
